#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Description:       :自定义dataloader
@Date     :2022/05/07 23:15:04
@Author      :Cosecant
@version      :1.0
                    _ooOoo_                     
                   o8888888o                    
                  88   .   88                    
                   (| -_- |)                    
                   O\  =  /O                    
                ____/`---'\____                 
              .'   \|     |/   `.               
             /   \|||  :  |||/   \              
            /  _||||| -:- |||||_  \             
            |   | \ \  -  /// |   |             
            | \_|  ''\---/''  |_/ |             
            \  .-\__  `-`  ___/-. /             
          ___`. .'  /--.--\  `. . __            
       .'' '<  `.___\_<|>_/___.'  >' ''.         
      | | :  `- \`.;`\ _ /`;.`/ - ` : | |       
      \  \ `-.   \_ __\ /__ _/   .-` /  /       
 ======`-.____`-.___\_____/___.-`____.-'======   
                    `=---='                     
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^   
        佛祖保佑        永无BUG                  
'''
from torch.utils.data import Dataset
import torch
import os
import scipy.io as sio
import numpy as np
import random
from matplotlib import pyplot as plt

def findAllMat(file_path):
    """
    @description  :找到file_path下所有的mat文件
    ---------
    @param  :fil_path :文件路径
    -------
    @Returns  :
    -------
    """
    for root, dirs, files in os.walk(file_path):
        for f in files:
            if (f.split('.')[-1] == "mat"):
                yield f


class MyDataSet(Dataset):
    def __init__(self, file_path, sequence_length ,dead_zone, step):
        super(MyDataSet, self).__init__()
        self.sequence_length=sequence_length
        self.step=step
        self.dead_zone=dead_zone
        self.matData = sio.loadmat(file_path)  #读取mat文件
        self.matData['emg']=self.matData['emg']-np.mean(self.matData['emg'], axis=0)
        self.matData['emg']=np.maximum(self.matData['emg'],-self.matData['emg'])
        
    def __getitem__(self, index):
        index2=int(index*self.step)
        #index2=index*(self.sequence_length-self.dead_zone)
        data = torch.Tensor(self.matData['emg'][index2:index2+self.sequence_length, ])
        label = torch.Tensor(self.matData['glove'][index2+self.sequence_length+self.dead_zone, ]) 
        """ data = torch.Tensor(self.matData['emg'][index:index+self.sequence_length, ])
        label = torch.Tensor(self.matData['glove'][index+self.sequence_length+self.dead_zone, ]) """
        return data, label

    def __len__(self):
        return int((len(self.matData['glove'])-self.sequence_length-self.dead_zone)/self.step)

class MyDataSet2(Dataset):
    def __init__(self, file_path, sequence_length ,dead_zone):
        super(MyDataSet2, self).__init__()
        self.sequence_length=sequence_length
        self.dead_zone=dead_zone
        self.matData = sio.loadmat(file_path)  #读取mat文件
        
    def __getitem__(self, index):
        index2=int(index*(self.sequence_length-self.dead_zone)/10)
        data = torch.Tensor(self.matData['emg'][index2:index2+self.sequence_length, ])
        label = torch.Tensor(self.matData['restimulus'][index2+self.sequence_length+self.dead_zone, ]) 
        """ data = torch.Tensor(self.matData['emg'][index:index+self.sequence_length, ])
        label = torch.LongTensor(self.matData['restimulus'][index+self.sequence_length+self.dead_zone, ]) """
        return data, label

    def __len__(self):
        return int((len(self.matData['restimulus'])-self.sequence_length-self.dead_zone)/(self.sequence_length-self.dead_zone)*10)
        #return int(len(self.matData['restimulus'])-self.sequence_length-self.dead_zone)

class MyDataSet3(Dataset):
    def __init__(self, file_path, sequence_length ,dead_zone):
        super(MyDataSet3, self).__init__()
        self.sequence_length=sequence_length
        self.dead_zone=dead_zone
        self.matData = sio.loadmat(file_path)  #读取mat文件
        self.change_index=[]
        mlist=['emg','glove','rerepetition','restimulus']
        for i in range(len(self.matData['restimulus'])-1):
            if(self.matData['restimulus'][i]!=self.matData['restimulus'][i+1]):
                self.change_index.append(i)
        middle_point=[] #因为变化总在0 *或* 0之间，所以记录* 0到 0 *变化的中点即为从动作0中点截取。
        middle_point.append(int(self.change_index.pop(0)/2)) #刚开始是从0 *开始的。
        last=int((len(self.matData['restimulus'])+self.change_index.pop())/2)
        while self.change_index:
            middle_point.append(int((self.change_index.pop(1)+self.change_index.pop(0))/2))
        middle_point.append(last)
        data_split=[]
        point1=middle_point.pop(0)
        while middle_point:
            point2=middle_point.pop(0)
            data_split.append([self.matData['emg'][point1:point2,],self.matData['restimulus'][point1:point2,]])
            point1=point2
        random.shuffle(data_split)
        self.shuffle_data={}
        foobar=data_split.pop()
        self.shuffle_data['emg']=foobar[0]
        self.shuffle_data['restimulus']=foobar[1]
        while data_split:
            foobar=data_split.pop()
            self.shuffle_data['emg']=np.append(self.shuffle_data['emg'], foobar[0], axis=0)
            self.shuffle_data['restimulus']=np.append(self.shuffle_data['restimulus'], foobar[1], axis=0)


    def __getitem__(self, index):
        index2=int(index*(self.sequence_length-self.dead_zone)/10)
        data = torch.Tensor(self.shuffle_data['emg'][index2:index2+self.sequence_length, ])
        label = torch.Tensor(self.shuffle_data['restimulus'][index2+self.sequence_length+self.dead_zone, ]) 
        """ data = torch.Tensor(self.matData['emg'][index:index+self.sequence_length, ])
        label = torch.LongTensor(self.matData['restimulus'][index+self.sequence_length+self.dead_zone, ]) """
        return data, label

    def __len__(self):
        return int((len(self.shuffle_data['restimulus'])-self.sequence_length-self.dead_zone)/(self.sequence_length-self.dead_zone)*10)
        #return int(len(self.matData['restimulus'])-self.sequence_length-self.dead_zone)

class MyDataSet4(Dataset):
    def __init__(self, matData, sequence_length , step):
        super(MyDataSet4, self).__init__()
        self.length=0
        self.step=step
        self.sequence_length=sequence_length
        self.matData=matData
        self.change_index=[]
        for i in range(len(self.matData['restimulus'])-1):
            if(self.matData['restimulus'][i]!=self.matData['restimulus'][i+1]):
                self.change_index.append(i)#找到所有的突变点
        #截掉突变的区域，截去前后共1秒 2000个点。
        self.split_point=[]
        self.split_point.append(500)
        while self.change_index:
            foo=self.change_index.pop(0)
            self.split_point.append(foo-1000)
            self.split_point.append(foo+1000)
        self.split_point.append(len(self.matData['restimulus'])-500)
        temp_window=[]
        for i in range(int(len(self.split_point)/2)): #split_point理论必是偶数
            point1=self.split_point[i*2]
            point2=self.split_point[i*2+1]
            temp_window.append([self.matData['emg'][point1:point2,],self.matData['restimulus'][point1:point2,]])
        temp_window = temp_window[1:-1:2]#只取动作区间
        #random.shuffle(temp_window)
        self.small_window=[]
        #savedata={}
        #savedata["emg"]=np.array(self.matData['emg'][0:1,])
        for i in range(len(temp_window)):
            window_num=int((len(temp_window[i][0])-sequence_length)/step)+1
            self.length+=window_num
            temp_window[i][0]=temp_window[i][0][0:(window_num-1)*step+sequence_length,]
            #savedata["emg"]=np.concatenate((savedata["emg"],temp_window[i][0]),0)
            temp_window[i][1]=temp_window[i][1][0:(window_num-1)*step+sequence_length,]
            for j in range(window_num):
                self.small_window.append([temp_window[i][0][j*step:j*step+sequence_length,],temp_window[i][1][j*step+sequence_length-1,]])
        #random.shuffle(self.small_window)
        #print(savedata["emg"].shape)
        #sio.savemat("test.mat", savedata)
        """ temp_list=[]
        print(temp_window[0][1].shape)
        for i in range(len(temp_window)):
            temp_list+=(list(np.array(temp_window[i][1]).ravel()))
        
        plt.plot(temp_list)
        plt.show() """

        """ self.shuffle_data={}
        foobar=temp_window.pop()
        self.shuffle_data['emg']=foobar[0]
        self.shuffle_data['restimulus']=foobar[1]
        while temp_window:
            foobar=temp_window.pop()
            self.shuffle_data['emg']=np.append(self.shuffle_data['emg'], foobar[0], axis=0)
            self.shuffle_data['restimulus']=np.append(self.shuffle_data['restimulus'], foobar[1], axis=0) """


    def __getitem__(self, index):
        #index2=int(index*self.step)
        data=torch.Tensor(self.small_window[index][0])
        label=torch.Tensor(self.small_window[index][1])
        """ data = torch.Tensor(self.shuffle_data['emg'][index2:index2+self.sequence_length, ])
        label = torch.Tensor(self.shuffle_data['restimulus'][index2+self.sequence_length, ])  """
        """ data = torch.Tensor(self.matData['emg'][index:index+self.sequence_length, ])
        label = torch.LongTensor(self.matData['restimulus'][index+self.sequence_length+self.dead_zone, ]) """
        return data, label

    def __len__(self):
        return self.length
        #return int(len(self.matData['restimulus'])-self.sequence_length-self.dead_zone)

class MyDataSet5(Dataset):
    def __init__(self, matData, sequence_length , step):
        super(MyDataSet5, self).__init__()
        self.length=0
        self.step=step
        self.sequence_length=sequence_length
        self.matData=matData
        self.change_index=[]
        for i in range(len(self.matData['restimulus'])-1):
            if(self.matData['restimulus'][i]!=self.matData['restimulus'][i+1]):
                self.change_index.append(i)#找到所有的突变点
        #截掉突变的区域，截去前后共1秒 2000个点。
        self.split_point=[]
        self.split_point.append(500)
        while self.change_index:
            foo=self.change_index.pop(0)
            self.split_point.append(foo-1000)
            self.split_point.append(foo+1000)
        self.split_point.append(len(self.matData['restimulus'])-500)
        temp_window=[]
        for i in range(int(len(self.split_point)/2)): #split_point理论必是偶数
            point1=self.split_point[i*2]
            point2=self.split_point[i*2+1]
            temp_window.append([self.matData['emg'][point1:point2,],self.matData['glove'][point1:point2,]])
        temp_window = temp_window[1:-1:2]#只取动作区间
        #random.shuffle(temp_window)
        self.small_window=[]
        #savedata={}
        #savedata["emg"]=np.array(self.matData['emg'][0:1,])
        for i in range(len(temp_window)):
            window_num=int((len(temp_window[i][0])-sequence_length)/step)+1
            self.length+=window_num
            temp_window[i][0]=temp_window[i][0][0:(window_num-1)*step+sequence_length,]
            #savedata["emg"]=np.concatenate((savedata["emg"],temp_window[i][0]),0)
            temp_window[i][1]=temp_window[i][1][0:(window_num-1)*step+sequence_length,]
            for j in range(window_num):
                self.small_window.append([temp_window[i][0][j*step:j*step+sequence_length,],temp_window[i][1][j*step:j*step+sequence_length,]])


    def __getitem__(self, index):
        #index2=int(index*self.step)
        data=torch.Tensor(self.small_window[index][0])
        label=torch.Tensor(self.small_window[index][1])
        """ data = torch.Tensor(self.shuffle_data['emg'][index2:index2+self.sequence_length, ])
        label = torch.Tensor(self.shuffle_data['restimulus'][index2+self.sequence_length, ])  """
        """ data = torch.Tensor(self.matData['emg'][index:index+self.sequence_length, ])
        label = torch.LongTensor(self.matData['restimulus'][index+self.sequence_length+self.dead_zone, ]) """
        return data, label

    def __len__(self):
        return self.length
if __name__ == '__main__':
    #fpath = './EMG_data/train/S1_E1_A1.mat'
    fpath = './EMG_data/test/S1_E1_A3.mat'
    matData = sio.loadmat(fpath)  #读取mat文件
    temp_dataset = MyDataSet4(matData, 2000, 100)

"""     for f in findAllMat(fpath):
        test=MyDataSet(fpath+f)
        mydata=test.__len__()
        print(mydata)
        break """